Won't it be a channels-first memory layout after the transpose?

The embeddings that come in are channel-first. I use the first transpose operation to make the embeddings channel-last. This helps with the padding, conv and batch_norm. After applying the operations I also use the next transpose to make the computeed embeddings channel-first again.

Am I missing something here?

Got it. I thought you were referring to the shape after the transpose.

I am assuming you've verified if this is indeed the sequence to follow when matching a torch.nn.Conv2d() with padding defined?

I have actually used this comment and colab notebook linked.

As per the official docs (https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html), these are the defaults:

• eps=1e-05, momentum=0.1

I have found out that in this repo all the BatchNorm layers had eps=1e-5 and momentum=0.9.

Also, true. The default momentum clearly isn't 0.9. @amyeroberts any inputs here?

But did changing the defaults help in the successful cross-loading by any chance?

It did not help me with anything in specific. I just did what all of the BatchNorm layers used.

What's the current issue again?

Would be better to have a detailed trace along with the steps to reproduce it somewhere.

The momentum parameter (rather confusingly) is different for PyTorch and TensorFlow. For momentum x in PyTorch, the equivalent value in TensorFlow is 1 - x.

My brain will remain dead for today.

Is there a reason why to have separate args and then kwargs? Usually, we only take kwargs here.

No specific reason. I can change it if you want!

No, it's like deviating from the standard implementations without reasoning. So.

from transformers import LevitFeatureExtractor, LevitModel, TFLevitModel
import torch
import tensorflow as tf
import numpy as np
from datasets import load_dataset

dataset = load_dataset("huggingface/cats-image")
image = dataset["test"]["image"][0]

feature_extractor = LevitFeatureExtractor.from_pretrained("facebook/levit-128S")

# Models
model_tf = TFLevitModel.from_pretrained("facebook/levit-128S", from_pt=True)
model_pt = LevitModel.from_pretrained("facebook/levit-128S")

# Inputs
inputs_tf = feature_extractor(image, return_tensors="tf")
inputs_pt = feature_extractor(image, return_tensors="pt")

# Outputs
outputs_tf = model_tf(**inputs_tf, training=False, output_hidden_states=False)
with torch.no_grad():
    outputs_pt = model_pt(**inputs_pt, output_hidden_states=False)

# Assertion
check_tf = outputs_tf.last_hidden_state
check_pt = outputs_pt.last_hidden_state
np.testing.assert_allclose(
    check_tf,
    check_pt,
    rtol=1e-5,
    atol=1e-5,
)

This creates a 100% mismatch.

And there's no warning when you're doing model_tf = TFLevitModel.from_pretrained("facebook/levit-128S", from_pt=True)? Like certain weights weren't used and so on?

All the weights are ported properly. The only thing that I see is with batch_norm parameters.

I have also followed this comment in order to see whether the parameters of the batch norm are important to be ported or not.

Hmm. If the weights have been cross-loaded successfully and the outputs are not being matched it means the intermediate computations are differing.

In this case, I would try to pick a few of the modules from the PT implementation, which could lead to differences in the outputs in the respective TF implementation.

Do you have any such modules in mind?

The TFLevitResidualLayer would be a good place to start with the inspection.

I am placing my bets on this module due to the use of random module.

If that's the case, I'd suggest starting your investigation from there. Specifically, ensure each intermediate operation in that module matches with one (PT) another (TF).